import os
import shap
import numpy as np
import pandas as pd
import utils.excel_util as excel_util
from sklearn.ensemble import RandomForestClassifier

# --------------------- 配置区 ---------------------
INPUT_XLSX = "./data/original.xlsx"
OUT_SHEET_NAME = "随机森林预测"  # 如果表名不同请修改
OUTPUT_DIR = "./out/analysis_outputs_v1.0"
READ_SHEET_NAME = "Sheet1"


# --------------------- 辅助函数 ---------------------
def experimental_stage(stage_type):
    stage_map = {"I": 40, "II": 20, "III": 20, "IV": 10}
    return stage_map.get(stage_type, 0)


def risk_level(risk_type):
    risk_map = {0: "低风险", 1: "高风险", -1: "未知风险"}
    return risk_map.get(risk_type, -1)


def check_risk_level(check_test):
    if check_test["标签_未来高风险"] == check_test["预测风险等级"]:
        return "true"
    else:
        return "false"


def calculate_feature_contributions(model, X, company_names, use_abs=True):
    """
    计算特征贡献度
    :param model: 模型
    :param X: 目标数据
    :param company_names: 机构名称
    :param use_abs: 取绝对值，消除负数，然后均值化到每个占比
    :return:
    """
    # 分离年份（如果有）
    year_col = None
    if "年份" in X.columns:
        year_col = X["年份"].values
        X_features = X.drop(columns=["年份"])
    else:
        X_features = X

    explainer = shap.TreeExplainer(model)
    shap_values = explainer.shap_values(X_features)

    # shap_values 可能是 list / ndarray
    if isinstance(shap_values, list):
        shap_values_class1 = shap_values[1]
    else:
        if shap_values.ndim == 3:  # (n_samples, n_features, n_classes)
            shap_values_class1 = shap_values[:, :, 1]  # 高风险类
        else:
            shap_values_class1 = shap_values  # 已经是 2D

    # 转 DataFrame（注意只用特征列）
    shap_df = pd.DataFrame(shap_values_class1, columns=X_features.columns)

    # ---- 计算占比 ----
    if use_abs:
        shap_for_ratio = shap_df.abs()
    else:
        shap_for_ratio = shap_df

    shap_df_ratio = shap_for_ratio.div(shap_for_ratio.sum(axis=1), axis=0)

    # 拼接企业名称和年份
    shap_df_ratio.insert(0, "企业名称", company_names.values)
    if year_col is not None:
        shap_df_ratio.insert(0, "年份", year_col)

    shap_df_ratio["预测分数"] = model.predict_proba(X_features)[:, 1]

    # base value
    base_value = explainer.expected_value
    if isinstance(base_value, (list, np.ndarray)):
        base_value = np.array(base_value)[1] if np.array(base_value).ndim > 0 else base_value

    return shap_df, shap_df_ratio, base_value


def to_long_for_tableau(shap_df_ratio_to_width):
    """
    适合Tableau的长表
    :param shap_df_ratio_to_width:
    :return:
    """
    # 保留的 id 列
    id_vars = ["年份", "企业名称", "预测分数"]
    # 特征列（除去 id_vars 的所有列）
    value_vars = [c for c in shap_df_ratio_to_width.columns if c not in id_vars]

    # 转换成长表
    radar_long = shap_df_ratio_to_width.melt(
        id_vars=id_vars,
        value_vars=value_vars,
        var_name="特征",
        value_name="占比"
    )
    return radar_long


def to_wide_for_excel(shap_df_ratio_to_width, index_name):
    """
    适合Excel的宽表
    :param shap_df_ratio_to_width
    :param index_name: 索引列名
    :return:
    """
    # 转换成长表
    radar_long = shap_df_ratio_to_width.pivot_table(
        index=index_name,
        columns="特征",
        values="占比",
        aggfunc="first"
    ).reset_index()

    return radar_long


# --------------------- 主流程 ---------------------
def main(input_xlsx=INPUT_XLSX, sheet_name=READ_SHEET_NAME, output_dir=OUTPUT_DIR):
    excel_data = excel_util.input_excel(input_xlsx, sheet_name)
    # 规范列名（去两端空格）
    excel_data.columns = [c.strip() for c in excel_data.columns]

    # ----------------- 案例一：随机森林算法模型 -----------------
    print("开始：案例一（随机森林算法模型）处理...")
    case1_cols = ["年份", "机构名称", "儿童试验项目数量", "当年项目数", "试验分期综合风险得分", "文件体系缺陷数",
                  "立项管理缺陷数", "质量管理缺陷数", "标签_未来高风险"]
    case1_exist = [c for c in case1_cols if c in excel_data.columns]
    case1 = excel_data[case1_exist].copy()
    case1["风险等级"] = case1["标签_未来高风险"].apply(risk_level)

    # 20222-2023 的数据进行训练，2024 的数据进测试
    train_data = case1[case1["年份"].between(2022, 2023)]
    test_data = case1[case1["年份"] == 2024]

    # ================================
    # 使用随机森林算法模型
    # ================================
    x = train_data[["儿童试验项目数量", "当年项目数", "试验分期综合风险得分", "文件体系缺陷数",
                    "立项管理缺陷数", "质量管理缺陷数"]]
    y = train_data["标签_未来高风险"]
    random_forest_model = RandomForestClassifier(n_estimators=200, oob_score=True, random_state=42)
    random_forest_model.fit(x, y)
    print("结束：案例一（随机森林算法模型）处理...")
    out_result = {}

    # OOB
    print("OOB score: {:.4f}".format(random_forest_model.oob_score_))
    out_result["OOB"] = pd.DataFrame({"OOB": random_forest_model.oob_score_}, index=[0])

    # 输出特征重要性
    importances = random_forest_model.feature_importances_
    feature_names = x.columns

    # 打印每个特征及其重要性
    for feature, importance in zip(feature_names, importances):
        print(f"{feature}: {importance:.4f}")
    # 输出特征值
    feature_importance_df = pd.DataFrame({
        '特征名称': feature_names,
        '重要性': importances
    })
    out_result["特征重要性"] = feature_importance_df

    # 提取机构名称列
    institution_names = case1["机构名称"].reset_index(drop=True)

    # 提取用于预测的特征列（与训练时一致）
    company_year_extract = case1[["儿童试验项目数量", "当年项目数", "试验分期综合风险得分",
                                  "文件体系缺陷数", "立项管理缺陷数", "质量管理缺陷数"]]
    year_extract = case1[["年份", "儿童试验项目数量", "当年项目数", "试验分期综合风险得分",
                          "文件体系缺陷数", "立项管理缺陷数", "质量管理缺陷数"]]

    # 使用模型进行预测
    predictions = random_forest_model.predict(company_year_extract)
    # 获取预测概率
    prediction_probs = random_forest_model.predict_proba(company_year_extract)[:, 1]

    # 将预测结果和分数转为 Series
    pred_series = pd.Series(predictions, name="预测结果")
    pred_prob_series = pd.Series(prediction_probs, name="预测分数")

    # 合并机构名称和预测结果
    result_df = pd.concat([institution_names, pred_series, pred_prob_series], axis=1)

    # 打印结果
    print("带机构名称的预测结果：")
    print(result_df)

    # 原始数据
    out_new_data = case1.copy()
    # 预测数据
    out_new_data["预测风险等级"] = result_df["预测结果"]
    out_new_data["预测分数"] = result_df["预测分数"]
    # 检测是否预测正确
    out_new_data["预测自检"] = out_new_data.apply(check_risk_level, axis=1)
    # 输出结果
    out_result["预测数据"] = out_new_data

    # 计算风险值，每个特征值的贡献（带正负的解释性表）
    shap_df_0, shap_df_ratio_0, base_value_0 = calculate_feature_contributions(
        random_forest_model,
        year_extract,
        institution_names,
        use_abs=False
    )
    # 保存结果
    out_result["特征贡献值-解释型"] = shap_df_0
    out_result["特征贡献占比-解释型"] = shap_df_ratio_0

    # 输出不带负数的占比（雷达图专用）
    shap_df_1, shap_df_ratio_1, base_value_1 = calculate_feature_contributions(
        random_forest_model,
        year_extract,
        institution_names,
        use_abs=True
    )

    # 保存结果
    out_result["特征贡献值-雷达图"] = shap_df_1
    out_result["特征贡献占比-雷达图"] = shap_df_ratio_1
    out_result["特征贡献占比-雷达图-转置"] = to_wide_for_excel(to_long_for_tableau(shap_df_ratio_1),
                                                                      ["年份", "企业名称", "预测分数"])

    excel_util.output_excel(output_dir, out_result, "case1_random_forest_all.xlsx")


if __name__ == "__main__":
    main()
